#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Datetime: 2020/9/15 11:10
# @Author  : CHEN Wang
# @Site    : 
# @File    : rebalance_date.py
# @Software: PyCharm 

"""
脚本说明: 再平衡日期模块，可支持多种频率，并且可以支持是每月月初或月末，每周几等等，同时支持是否需要为交易日，方便测试日历路径效应
"""

import pandas as pd
from quant_researcher.quant.project_tool.db_operator import db_conn
from quant_researcher.quant.project_tool import time_tool
from quant_researcher.quant.project_tool.logger.my_logger import LOG
from quant_researcher.quant.datasource_fetch.common_data_api import t_trade_date


def get_custom_freq(freq):
    b_cal = t_trade_date.get_trade_date_calendar()
    if freq == 'month':
        freq = pd.offsets.CBMonthEnd(calendar=b_cal)
    elif freq == 'week':
        freq = 'W-FRI'

    return freq


def get_rebalance_date(start_date=None, end_date=None, freq='M', period=1, business_day=False,
                   **kwargs):
    """
    再平衡日期模块，可支持多种频率，并且可以支持是每月月初或月末，每周几等等，同时支持是否需要为交易日，方便测试日历路径效应

    :param str, start_date: 策略开始时间， '2020-01-01'
    :param str, end_date: 策略结束时间， '2020-08-01'， 未输入则默认当前最新时间
    :param str, freq: 再平衡频率（调仓频率）支持pandas提供的频率，‘M’频率为月末调仓，‘MS’频率为月初调仓，'W-FRI'为周五调仓，
        其他频率详情参见 https://pandas.pydata.org/pandas-docs/stable/user_guide/timeseries.html#offset-aliases
    :param int, period: 再平衡周期，是每几个月，或者每几周
    :param bool, business_day: 再平衡是否需要是交易日， 默认为False
    :param kwargs:
        -str, market_code: 交易所代码，默认为'0101001',即A股交易日历
        -bool, initial_start: 起始日期是否需要包含在再平衡日期内， 默认为True
        -bool, initial_end: 结束日期是否需要包含在再平衡日期内， 默认为False
    :return: 再平衡日期列表， 每个日期的格式为'2020-01-01'
    """

    if start_date is None:
        LOG.warning(f"start_date 必须输入")
    if end_date is None:
        end_date = time_tool.get_today(marker='with_n_dash')

    initial_start = kwargs.pop('initial_start', True)  # 默认第一个日期为start_date, 或start_date起的最早一个交易日
    initial_end = kwargs.pop('initial_end', False)  # 默认最后日期为end_date, 或end_date往前最后一个交易日

    rebalance_freq = str(period) + str(freq)
    rebalance_date = pd.date_range(start_date, end_date, freq=rebalance_freq).tolist()
    rebalance_date = [x.strftime("%Y-%m-%d") for x in rebalance_date]

    if initial_start:  # 添加start_date到再平衡日期列表中
        if rebalance_date[0] != start_date:
            rebalance_date = [start_date] + rebalance_date
    if initial_end:  # 添加end_date到再平衡日期列表中
        if rebalance_date[-1] != end_date:
            rebalance_date = rebalance_date + [end_date]

    if business_day:  # 要求再平衡日期必须为交易日
        market_code = kwargs.pop('market_code', '0101001')  # 默认使用沪深交易所日历
        if freq[0] == 'M':  # 即再平衡以月为单位
            for i, date in enumerate(rebalance_date):
                if t_trade_date.is_a_trade_date(the_day=date, market_code=market_code):
                    continue
                else:  # 如果该平衡日期为非交易日
                    if freq[-1] == 'S':  # 再平衡为月初
                        temp = t_trade_date.date_shifter_by_trade_date(date, how_many=1, before_or_after='after')
                    if freq[-1] == 'M':  # 再平衡为月末
                        temp = t_trade_date.date_shifter_by_trade_date(date, how_many=1, before_or_after='before')
                    rebalance_date[i] = temp

        if freq[0] == 'W':  # 即再平衡以周为单位
            for i, date in enumerate(rebalance_date):
                if t_trade_date.is_a_trade_date(the_day=date, market_code=market_code):
                    continue
                else:  # 如果该平衡日期为非交易日
                    if freq[-3:] == 'FRI':  # 再平衡为周五
                        temp = t_trade_date.date_shifter_by_trade_date(date, how_many=1, before_or_after='before')
                    else:  # 再平衡为周内其他
                        temp = t_trade_date.date_shifter_by_trade_date(date, how_many=1, before_or_after='after')
                    rebalance_date[i] = temp

    # 去重并保持原有顺序
    new_rebalance_date = list(set(rebalance_date))
    new_rebalance_date.sort(key=rebalance_date.index)

    return new_rebalance_date


def multi_period_process(inner_func, only_latest=True, **kwargs):
    """
    多期处理函数

    :param inner_func: 每期处理的函数
    :param only_latest: 是否只返回最新一期
    :param kwargs: 如果不只是获取最新一期，则需要传入开始时间，及更新频率
        - str start_date: 开始时间 "2020-01-01"
        - str end_date: 开始时间 "2020-09-01"
        - str freq: 更新频率， 支持'D', 'W', 'M', 'Q', 其他频率
        - int period: 再平衡周期，是每几个月，或者每几周
        详情参见 https://pandas.pydata.org/pandas-docs/stable/user_guide/timeseries.html#offset-aliases
    :return: 只返回最新一期为dataframe，返回多期则为以日期为key，当期基金dataframe为值的字典
    """

    func_kwargs = kwargs.pop('func_kwargs', {})
    end_date = kwargs.pop('end_date', None)
    if end_date is None:
        end_date = time_tool.get_today(marker='with_n_dash')
    if only_latest:
        return inner_func(end_date, **func_kwargs)
    else:
        assert all([x in kwargs.keys() for x in ['start_date', 'freq']]), '请确认输入start_date和freq参数'
        start_date = kwargs.pop('start_date')
        freq = kwargs.pop('freq', 'M')
        period = kwargs.pop('period', 1)  # 周期默认为1
        date_range = get_rebalance_date(start_date, end_date, freq=freq, period=period, business_day=True, initial_start=True, initial_end=False)
        date_list = []
        for date in date_range:
            print(f'开始处理{date}')
            temp = inner_func(date, **func_kwargs)
            temp['date'] = date
            date_list.append(temp)
        date_funds_df = pd.concat(date_list)

        return date_funds_df


if __name__ == '__main__':
    start_date = '2020-01-31'
    end_date = '2020-08-31'
    rebalance_date = get_rebalance_date(start_date, end_date, freq='M', period=1, business_day=True,  initial_start=True, initial_end=True)

    # 测试multi_period_filter_process
    def inner(day):
        conn = db_conn.get_derivative_data_conn()
        ganggu_df = pd.read_sql(f"select fund_code, fund_sname "
                                f"from mf_bd_fndinfo "
                                f"where ((fund_state = 0) or (fund_enddate >='{day}')) "
                                f"and (fund_sname like '%%香港%%' or fund_sname like '%%恒生%%')", conn)
        return ganggu_df
    start_date = '2020-01-31'
    end_date = '2020-08-31'
    rebalance_date = multi_period_process(inner_func=inner, only_latest=False, start_date=start_date, end_date=end_date, freq='M', period=1)


